;;;

;;;; eigen.lisp
;;;;
;;;; Support routines for eigenvectors and eigenvalues
;;;;


(in-package :lib.gl)

;;; ---- Householder Transformation ----
;;;             v*(transpose(v))
;;; H = I - 2 * ----------------
;;;              magnitude(v)^2
;;; If v is normalized, i.e. u = v/magnitude(v)
;;;
;;; H = I - 2*u*transpose(u)
;;;

(defun householder-transform (v)
  (let ((size (vector-length v))
	(u (m/ v (magnitude v))))
    (m- (identity-matrix size)
	(m* 2 u (transpose u)))))

(defun embed-householder (h rank)
  (let* ((s (matrix-rows h))
	 (di (- rank s))
	 (u (make-matrix rank rank)))
    (dotimes (i di)
      (setf (mref u i i) 1.0))
    (dotimes (i s u)
      (dotimes (j s)
	(setf (mref u (+ i di) (+ j di)) (mref h i j))))))

(defun sub-vector (v start)
  (apply #'vector
	 (iter (for i from start below (vector-length v))
	       (collecting (mref v i)))))

(defun compute-householder (a col)
  (let* ((x (sub-vector (get-column a col) (1+ col)))
	 (theta (magnitude x))
	 (z (zero-vector (vector-length x))))
    (setf (mref z 0) 1.0)
    (embed-householder (householder-transform (m+ x (m* theta z)))
		       (matrix-cols a))))

(defun %tridiagonalize (a n)
  (let* ((r (matrix-cols a))
	 (u (compute-householder a (- r n))))
    (cond ((= (1- n) 2) (m* u a u))
	  (t (%tridiagonalize (M* u a u) (1- n))))))

(defun hessenberg-form (a)
  (let ((r (matrix-rows a)))
    (%tridiagonalize a r)))

(defun shifted-qr-step (a)
  (let* ((nr (matrix-rows a))
	 (alpha_k (mref a (1- nr) (1- nr))))
    (multiple-value-bind (q r) (qr-factorization (m- a (m* alpha_k (identity-matrix nr))))
      (m+ (m* r q) (m* alpha_k (identity-matrix nr))))))

(defun unshifted-qr-step (a)
  (multiple-value-bind (q r) (qr-factorization a)
    (m* r q)))

(defun run-qr-step (a tolerance &optional (method #'shifted-qr-step))
  (let ((n (1- (matrix-rows a))))
    (cond ((< (abs (mref a n (1- n))) tolerance) a)
	   (t (run-qr-step (funcall method a) tolerance method)))))

;;; Test for diagonal matrices and save yourself some time and effort!
(defun find-eigenvalues (a &key (tolerance 0.001) (method #'shifted-qr-step))
  (let* ((res (run-qr-step a tolerance method))
	 (n (matrix-rows a))
	 (ev (mref res (1- n) (1- n))))
    (cond ((<= n 2) (list ev (mref res 0 0)))
	  (t (let ((res (find-eigenvalues (submatrix res 0 0 (- n 2) (- n 2)) :tolerance tolerance :method method)))
	       (push ev res))))))

(defun find-eigenvectors (a eigenvalues)
  (if (diagonal-p a)
      (iter (for i from 0 below (matrix-cols a))
	    (for v = (zero-vector 3))
	    (setf (mref v i) 1.0)
	    (collect v))
      (iter (for e in eigenvalues)
	    (collect (inverse-power-step a e)))))


(defun inverse-power-step (a eigenvalue &optional (normalize t))
  (let* ((n (matrix-rows a))
	 (ashift (m- a (m* eigenvalue (identity-matrix n)))))
    (multiple-value-bind (l u) (doolittle-lu-decomposition ashift)
      (let* ((lx0 (gaussian-forward-elimination (augment-matrix l (ones-vector n))))
	     (ux0 (gaussian-back-substitution (augment-matrix u (get-column lx0 n))))
	     (eigvec (get-column ux0 n)))
	(if normalize
	    (m/ eigvec (magnitude eigvec))
	    eigvec)))))

;;;; -------- ;;;;

(defun test-hessenberg-1 ()
  (assert (zero-matrix-p (m- (hessenberg-form (%matrix '((1.0 0.0 1.0)
							 (0.0 1.0 1.0)
							 (1.0 1.0 0.0))))
			     (%matrix '((1.0 -1.0 0.0)
					(-1.0 0.0 1.0)
					(0.0 1.0 1.0)))))))

(defun test-hessenberg-2 ()
  (let ((a (%matrix '((4.0  1.0 -2.0  2.0)
		      (1.0  2.0  0.0  1.0)
		      (-2.0 0.0  3.0 -2.0)
		      (2.0  1.0 -2.0 -1.0))))
	(r (%matrix `((4.0 -3.0 0.0 0.0)
		      (-3.0 ,(/ 10 3.0) ,(/ -5 3.0) 0.0)
		      (0.0 ,(/ -5 3.0) ,(/ -33 25.0) ,(/ 68 75.0))
		      (0.0 0.0 ,(/ 68 75.0) ,(/ 149 75.0))))))
    (zero-matrix-p (m- (hessenberg-form a) r))
    (hessenberg-form a)))

(defun test-qr-1 ()
  (let ((a (%matrix '((-149.0 -50.0 -154.0)
		      (537.0 180.0 546.0)
		      (-27.0 -9.0 -25.0)))))
    (run-qr-step a 0.0001)))

(defun qr-benchmark ()
  (let ((a (%matrix '((-149.0 -50.0 -154.0)
		      (537.0 180.0 546.0)
		      (-27.0 -9.0 -25.0)))))
    (time (dotimes (i 100) (run-qr-step a 0.0001 #'shifted-qr-step)))
    (time (dotimes (i 100) (run-qr-step a 0.0001 #'unshifted-qr-step)))))
	  

(defun test-eigenvalues-1 ()
  (let ((a (%matrix '((-149.0 -50.0 -154.0)
		      (537.0 180.0 546.0)
		      (-27.0 -9.0 -25.0)))))
    (find-eigenvalues a :tolerance 0.00001))) ;; Should return close to (2 1 3)

(defun test-tough-eigenvalues ()
  (let ((a (%matrix '((0.0 2.0 0.0 -1.0)
		      (1.0 0.0 0.0 0.0)
		      (0.0 1.0 0.0 0.0)
		      (0.0 0.0 1.0 0.0)))))
    (find-eigenvalues a :tolerance 0.01))) ;; Yep, it doesn't work.